import socket
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from config_loader import ConfigLoader
from ssh_connect import SSHConnect

class CommandDistributor:
    def __init__(self, config: ConfigLoader):
        self.config = config
        self.connect_pool = {}
        self.lock = threading.Lock()

    def test_connective(self):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(15)
        servers = self.config.get_connectInfo()
        for name, info in servers.items():  # type: ignore
            print(f"Testing {name} is connective:")
            host, port, username, pwd, cert_file = info.getLoginParam()
            try:
                result = sock.connect_ex((host, int(port)))
                print(f"Connection result for {name}: {result}")
            except Exception as e:
                print(f"Error connecting to {name}: {e}")

    def connect_all(self):
        servers = self.config.get_connectInfo()
        if not isinstance(servers, dict):
            print("ConfigLoader.get_connectInfo() did not return a dictionary of servers.")
            return
        for name, info in servers.items():
            print(f"trying to connect {name}")
            ssh = SSHConnect(info)
            ssh.connect()  # Actually connect here
            if not ssh.status:
                continue
            self.connect_pool[name] = ssh
        if len(self.connect_pool) == 0:
            print("No any successful connection")
            raise ConnectionError

    def distributed_command(self, command: str):
        """Execute command on all servers in parallel and collect results"""
        if not self.connect_pool:
            print("No connections available. Please run connect_all() first.")
            return {}
        
        results = {}
        
        def _execute_command(name, ssh_connect, cmd):
            """Helper function to execute command and return result"""
            try:
                print(f"Executing command on {name}: {cmd}")
                success = ssh_connect.command(cmd)
                return name, success, "Command executed successfully" if success else "Command failed"
            except Exception as e:
                return name, False, f"Error: {str(e)}"
        
        with ThreadPoolExecutor(max_workers=len(self.connect_pool)) as executor:
            # Submit all tasks
            future_to_name = {
                executor.submit(_execute_command, name, ssh_connect, command): name 
                for name, ssh_connect in self.connect_pool.items()
            }
            
            # Collect results as they complete
            for future in as_completed(future_to_name):
                name = future_to_name[future]
                try:
                    server_name, success, message = future.result()
                    results[server_name] = {
                        'success': success,
                        'message': message
                    }
                    status = "✓" if success else "✗"
                    print(f"{status} {server_name}: {message}")
                except Exception as e:
                    results[name] = {
                        'success': False,
                        'message': f"Execution failed: {str(e)}"
                    }
                    print(f"✗ {name}: Execution failed - {str(e)}")
        
        return results
    
    def fileupload(self, inpath: str, distpath: str):
        if not self.connect_pool:
            print("No connections available. Please run connect_all() first.")
            return {}
        
        filename = inpath.split("/")[-1]
        distpath = "".join([distpath, filename])
        
        results = {}

        def _execute_command(name, ssh_connect: SSHConnect, inpath: str, distpath: str):
            """Helper function to execute command and return result"""
            try:
                print(f"Upload file to {distpath} on {name}")
                success = ssh_connect.fileupload(inpath, distpath)
                return name, success, "Command executed successfully" if success else "Command failed"
            except Exception as e:
                return name, False, f"Error: {str(e)}"
        
        with ThreadPoolExecutor(max_workers=len(self.connect_pool)) as executor:
            # Submit all tasks
            future_to_name = {
                executor.submit(_execute_command, name, ssh_connect, inpath, distpath): name 
                for name, ssh_connect in self.connect_pool.items()
            }
            
            # Collect results as they complete
            for future in as_completed(future_to_name):
                name = future_to_name[future]
                try:
                    server_name, success, message = future.result()
                    results[server_name] = {
                        'success': success,
                        'message': message
                    }
                    status = "✓" if success else "✗"
                    print(f"{status} {server_name}: {message}")
                except Exception as e:
                    results[name] = {
                        'success': False,
                        'message': f"Execution failed: {str(e)}"
                    }
                    print(f"✗ {name}: Execution failed - {str(e)}")
        
        return results

    def close_all_connections(self):
        """Close all SSH connections"""
        for name, ssh_connect in self.connect_pool.items():
            try:
                ssh_connect.close()
                print(f"Closed connection to {name}")
            except Exception as e:
                print(f"Error closing connection to {name}: {e}")
        self.connect_pool.clear()